package org.example.aitest.filter;

import org.springframework.util.StreamUtils;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;
import java.nio.charset.StandardCharsets;

/**
 * XSS过滤处理

 */
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper
{
    private final byte[] cachedBody;

    public XssHttpServletRequestWrapper(HttpServletRequest request) throws IOException {
        super(request);
//        this.cachedBody = StreamUtils.copyToByteArray(request.getInputStream());
        try {
            InputStream is = request.getInputStream();
            this.cachedBody = is != null ? StreamUtils.copyToByteArray(is) : new byte[0];
        }finally {
            request.getInputStream().close();
        }
    }

    // 新增方法：用于调试获取缓存内容
    public byte[] getCachedBody() {
        return this.cachedBody.clone();
    }

    /**
     * @param request
     */
    public XssHttpServletRequestWrapper(HttpServletRequest request, byte[] cachedBody)
    {
        super(request);
        this.cachedBody = cachedBody;
    }

    @Override
    public String[] getParameterValues(String name)
    {
        String[] values = super.getParameterValues(name);
        if (values != null)
        {
            int length = values.length;
            String[] escapesValues = new String[length];
            return escapesValues;
        }
        return super.getParameterValues(name);
    }

    @Override
    public ServletInputStream getInputStream() throws IOException
    {
        // 非json类型，直接返回
        /*if (!isJsonRequest())
        {
            System.out.println("非json");
            return super.getInputStream();
        }

        // 为空，直接返回
        String json = IOUtils.toString(super.getInputStream(), "utf-8");
        if (StringUtils.isEmpty(json))
        {
            System.out.println("空json");
            return super.getInputStream();
        }*/

        return new CachedServletInputStream(this.cachedBody.clone());
        // xss过滤
        /*json = EscapeUtil.clean(json).trim();
        byte[] jsonBytes = json.getBytes("utf-8");
        final ByteArrayInputStream bis = new ByteArrayInputStream(jsonBytes);
        return new ServletInputStream()
        {
            @Override
            public boolean isFinished()
            {
                return true;
            }

            @Override
            public boolean isReady()
            {
                return true;
            }

            @Override
            public int available() throws IOException
            {
                return jsonBytes.length;
            }

            @Override
            public void setReadListener(ReadListener readListener)
            {
            }

            @Override
            public int read() throws IOException
            {
                return bis.read();
            }
        };*/
    }

    /**
     * 是否是Json请求
     * 
     * @param request
     */
    public boolean isJsonRequest()
    {
        String contentType = this.getContentType();
        return contentType != null && contentType.toLowerCase().startsWith("application/json");
//        String header = super.getHeader(HttpHeaders.CONTENT_TYPE);
//        return StringUtils.startsWithIgnoreCase(header, MediaType.APPLICATION_JSON_VALUE);
    }

    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(
                this.getInputStream(), StandardCharsets.UTF_8));
    }

    // 自定义 ServletInputStream
    static class CachedServletInputStream extends ServletInputStream {
        private final ByteArrayInputStream buffer;

        public CachedServletInputStream(byte[] body) {
            this.buffer = new ByteArrayInputStream(body);
        }

        @Override public int read() { return buffer.read(); }
        @Override public boolean isFinished() { return buffer.available() == 0; }
        @Override public boolean isReady() { return true; }
        @Override public void setReadListener(ReadListener listener) {
            throw new UnsupportedOperationException();
        }
    }
}